import pandas as pd
from PySide6.QtCore import QObject, Signal
from typing import Dict, List, Optional
from utils.db_utils import DatabaseManager

class DataModel(QObject):
    """数据模型类，负责数据的加载、处理和转换"""
    data_loaded = Signal(str, pd.DataFrame)  # 信号：文件名和数据
    data_updated = Signal(str, pd.DataFrame)  # 信号：数据更新

    def __init__(self):
        super().__init__()
        self._data_cache: Dict[str, pd.DataFrame] = {}
        self._recent_files: List[str] = []
        self.db_manager = DatabaseManager()
        self._init_database()

    def _init_database(self):
        """初始化数据库"""
        self.db_manager.create_tables()
        self.db_manager.insert_sample_data()

    def load_data(self, well_name: str) -> Optional[pd.DataFrame]:
        """从数据库加载井数据"""
        try:
            # 查询井数据
            query = """
            SELECT depth, value 
            FROM well_data 
            WHERE well_name = %s 
            ORDER BY depth
            """
            results = self.db_manager.execute_query(query, (well_name,))
            
            if results:
                # 转换为DataFrame
                df = pd.DataFrame(results)
                self._data_cache[well_name] = df
                self._update_recent_files(well_name)
                self.data_loaded.emit(well_name, df)
                return df
            return None
        except Exception as e:
            print(f"Error loading data: {e}")
            return None

    def get_well_list(self) -> List[Dict]:
        """获取井列表"""
        query = "SELECT well_name, location, depth, description FROM well_info"
        return self.db_manager.execute_query(query) or []

    def get_data(self, well_name: str) -> Optional[pd.DataFrame]:
        """获取已加载的数据"""
        if well_name in self._data_cache:
            return self._data_cache[well_name]
        return self.load_data(well_name)

    def update_data(self, well_name: str, data: pd.DataFrame):
        """更新数据"""
        self._data_cache[well_name] = data
        self.data_updated.emit(well_name, data)

    def _update_recent_files(self, well_name: str):
        """更新最近文件列表"""
        if well_name in self._recent_files:
            self._recent_files.remove(well_name)
        self._recent_files.insert(0, well_name)
        self._recent_files = self._recent_files[:10]  # 保持最近10个文件

    def get_recent_files(self) -> List[str]:
        """获取最近文件列表"""
        return self._recent_files

    def close(self):
        """关闭数据库连接"""
        self.db_manager.close() 